from .ACDC import get_acdc_loader, get_ssl_acdc_loader
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


cifar10_mean = (0.4914, 0.4822, 0.4465)
cifar10_std = (0.2471, 0.2435, 0.2616)
cifar100_mean = (0.5071, 0.4867, 0.4408)
cifar100_std = (0.2675, 0.2565, 0.2761)
normal_mean = (0.5, 0.5, 0.5)
normal_std = (0.5, 0.5, 0.5)

def build_loader(args):
    if args.datasets == "acdc":
        label_loader, unlabel_loader, test_loader = get_ssl_acdc_loader(
            root=args.data_path,
            train_crop_size=args.train_crop_size,
            batch_size=args.batch_size,
            unlabel_batch_size=args.unlabel_batch_size,
            label_num=args.label_num)
        return label_loader, unlabel_loader, test_loader
    
    elif args.datasets=="cifar10":


        train_transform = transforms.Compose([
            # transforms.RandomResizedCrop(32, scale=(0.75, 1.0), ratio=(1.0, 1.0)),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ToTensor(),
            transforms.Normalize(mean=cifar10_mean, std=cifar10_std),
        ])

        test_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=cifar10_mean, std=cifar10_std),
        ])
            
        train_dataset = datasets.CIFAR10(
            root=args.data_path,
            train=True,
            download=True,
            transform=train_transform,
        )

        test_dataset = datasets.CIFAR10(
            root=args.data_path,
            train=False,
            download=True,
            transform=test_transform,
        )
        train_loader =DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True,num_workers=4)
        test_loader =DataLoader(test_dataset, batch_size=args.batch_size,shuffle=False,num_workers=4)

        return train_loader,test_loader


    else:
        raise NotImplementedError
